{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Enhanced Swarm Orchestration with AG2 \n",
    "\n",
    "AG2's swarm orchestration provides a flexible and powerful method of managing a conversation with multiple agents, tools, and transitions.\n",
    "\n",
    "In this notebook, we look at more advanced features of the swarm orchestration.\n",
    "\n",
    "If you are new to swarm, check out [this notebook](https://docs.ag2.ai/latest/docs/use-cases/notebooks/notebooks/agentchat_swarm), where we introduce the core features of swarms including global context variables, hand offs, and initiating a swarm chat.\n",
    "\n",
    "In this notebook we're going to demonstrate these features AG2's swarm orchestration:\n",
    "\n",
    "- Updating an agent's state\n",
    "- Conditional handoffs\n",
    "- Nested chats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{=mdx}\n",
    ":::info Requirements\n",
    "Install `ag2`:\n",
    "```bash\n",
    "pip install -U ag2[openai]\n",
    "```\n",
    "\n",
    "> **Note:** If you have been using `autogen` or `ag2`, all you need to do is upgrade it using:  \n",
    "> ```bash\n",
    "> pip install -U autogen\n",
    "> ```\n",
    "> or  \n",
    "> ```bash\n",
    "> pip install -U ag2\n",
    "> ```\n",
    "> as `autogen`, and `ag2` are aliases for the same PyPI package.  \n",
    "\n",
    "\n",
    "For more information, please refer to the [installation guide](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/installing-ag2).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demonstration\n",
    "\n",
    "We're creating this customer service workflow for an e-commerce platform. Customers can ask about the status of their orders, but they must be authenticated to do so.\n",
    "\n",
    "![Swarm Enhanced Demonstration](https://media.githubusercontent.com/media/ag2ai/ag2/refs/heads/main/notebook/swarm_enhanced_01.png)\n",
    "\n",
    "Key aspects of this swarm are:\n",
    "\n",
    "1. System messages are customised, incorporating the context of the workflow\n",
    "2. Handoffs are conditional, only being available when they are relevant\n",
    "3. A nested chat handles the order retrieval and summarisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# IMPORTS\n",
    "import json\n",
    "from typing import Any\n",
    "\n",
    "# AG2 imports\n",
    "from autogen import (\n",
    "    AfterWork,\n",
    "    AfterWorkOption,\n",
    "    ContextExpression,\n",
    "    ConversableAgent,\n",
    "    LLMConfig,\n",
    "    OnCondition,\n",
    "    OnContextCondition,\n",
    "    SwarmResult,\n",
    "    UpdateSystemMessage,\n",
    "    UserProxyAgent,\n",
    "    initiate_swarm_chat,\n",
    "    register_hand_off,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set your API Endpoint\n",
    "\n",
    "The [`LLMConfig.from_json`](https://docs.ag2.ai/latest/docs/api-reference/autogen/llm_config/LLMConfig) method loads a list of configurations from an environment variable or a json file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load configuration from an OAI_CONFIG_LIST file\n",
    "# and filter it to OpenAI's gpt-5\n",
    "# Create an LLM configuration with our config\n",
    "llm_config = LLMConfig.from_json(\n",
    "    path=\"OAI_CONFIG_LIST\",\n",
    "    cache_seed=42,\n",
    "    temperature=1,\n",
    "    timeout=120,  # change the cache_seed for different trials\n",
    ").where(tags=[\"gpt-5\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Context used for controlling the workflow and injecting information\n",
    "# into some agent's system messages\n",
    "workflow_context = {\n",
    "    # customer details\n",
    "    \"customer_name\": None,  # Injected into Triage & Order Management Agent's system messages\n",
    "    \"logged_in_username\": None,  # Helps validate order\n",
    "    # workflow status\n",
    "    \"logged_in\": False,  # Control the flow\n",
    "    # order enquiry details\n",
    "    \"has_order_id\": False,  # Control the flow\n",
    "    \"order_id\": None,  # Order validation and included in system messages\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Databases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mock Databases\n",
    "\n",
    "# User Database with the username as the key\n",
    "USER_DATABASE = {\n",
    "    \"mark\": {\n",
    "        \"full_name\": \"Mark Sze\",\n",
    "    },\n",
    "    \"kevin\": {\n",
    "        \"full_name\": \"Yiran Wu\",\n",
    "    },\n",
    "}\n",
    "\n",
    "# Order Database with the order number as the key\n",
    "# The user field aligns with the username in USER_DATABASE\n",
    "ORDER_DATABASE = {\n",
    "    \"TR13845\": {\n",
    "        \"user\": \"mark\",\n",
    "        \"order_number\": \"TR13845\",\n",
    "        # order status: order_received, shipped, delivered, return_started, returned\n",
    "        \"status\": \"shipped\",\n",
    "        # return status: N/A, return_started, return_shipped, return_delivered, refund_issued\n",
    "        \"return_status\": \"N/A\",\n",
    "        \"product\": \"matress\",\n",
    "        \"link\": \"https://www.example.com/TR13845\",\n",
    "        \"shipping_address\": \"123 Main St, State College, PA 12345\",\n",
    "    },\n",
    "    \"TR14234\": {\n",
    "        \"user\": \"kevin\",\n",
    "        \"order_number\": \"TR14234\",\n",
    "        \"status\": \"delivered\",\n",
    "        \"return_status\": \"N/A\",\n",
    "        \"product\": \"pillow\",\n",
    "        \"link\": \"https://www.example.com/TR14234\",\n",
    "        \"shipping_address\": \"123 Main St, State College, PA 12345\",\n",
    "    },\n",
    "    \"TR29384\": {\n",
    "        \"user\": \"mark\",\n",
    "        \"order_number\": \"TR29384\",\n",
    "        \"status\": \"delivered\",\n",
    "        \"return_status\": \"N/A\",\n",
    "        \"product\": \"bed frame\",\n",
    "        \"link\": \"https://www.example.com/TR29384\",\n",
    "        \"shipping_address\": \"123 Main St, State College, PA 12345\",\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agent's Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ORDER FUNCTIONS\n",
    "def check_order_id(order_id: str, context_variables: dict) -> SwarmResult:\n",
    "    \"\"\"Check if the order ID is valid\"\"\"\n",
    "    order_id = order_id.upper().strip()\n",
    "\n",
    "    # Restricts order to checking to the logged in user\n",
    "    if (\n",
    "        context_variables[\"logged_in_username\"]\n",
    "        and order_id in ORDER_DATABASE\n",
    "        and ORDER_DATABASE[order_id][\"user\"] == context_variables[\"logged_in_username\"]\n",
    "    ):\n",
    "        return SwarmResult(\n",
    "            context_variables=context_variables, values=f\"Order ID {order_id} is valid.\", agent=order_triage_agent\n",
    "        )\n",
    "    else:\n",
    "        # Order ID wasn't found or doesn't belong to the logged in user\n",
    "        return SwarmResult(\n",
    "            context_variables=context_variables,\n",
    "            values=f\"Order ID {order_id} is invalid. Please ask for the correct order ID.\",\n",
    "            agent=order_triage_agent,\n",
    "        )\n",
    "\n",
    "\n",
    "def record_order_id(order_id: str, context_variables: dict) -> SwarmResult:\n",
    "    \"\"\"Record the order ID in the workflow context\"\"\"\n",
    "    order_id = order_id.upper().strip()\n",
    "\n",
    "    if order_id not in ORDER_DATABASE:\n",
    "        # Order ID wasn't found or doesn't belong to the logged in user\n",
    "        # Transfer back to the order triage agent\n",
    "        return SwarmResult(\n",
    "            context_variables=context_variables,\n",
    "            values=f\"Order ID {order_id} not found. Please ask for the correct order ID.\",\n",
    "            agent=order_triage_agent,\n",
    "        )\n",
    "\n",
    "    context_variables[\"order_id\"] = order_id\n",
    "    context_variables[\"has_order_id\"] = True\n",
    "\n",
    "    # Order ID has been recorded, transfer to the order management agent\n",
    "    return SwarmResult(\n",
    "        context_variables=context_variables, values=f\"Order ID Recorded: {order_id}\", agent=order_mgmt_agent\n",
    "    )\n",
    "\n",
    "\n",
    "# AUTHENTICATION FUNCTIONS\n",
    "def login_customer_by_username(username: str, context_variables: dict) -> SwarmResult:\n",
    "    \"\"\"Get and log the customer in by their username\"\"\"\n",
    "    username = username.lower().strip()\n",
    "\n",
    "    if username in USER_DATABASE:\n",
    "        context_variables[\"customer_name\"] = USER_DATABASE[username][\"full_name\"]\n",
    "        context_variables[\"logged_in_username\"] = username\n",
    "        context_variables[\"logged_in\"] = True\n",
    "\n",
    "        # Successful login, transfer to the order triage agent\n",
    "        return SwarmResult(\n",
    "            context_variables=context_variables,\n",
    "            values=f\"Welcome back our customer, {context_variables['customer_name']}! Please continue helping them.\",\n",
    "            # Note: We aren't specifying an agent here as the OnContextCondition handoff the Authentication Agent will handle it\n",
    "        )\n",
    "    else:\n",
    "        # User not found, transfer back to the authentication agent\n",
    "        return SwarmResult(\n",
    "            context_variables=context_variables,\n",
    "            values=f\"User {username} not found. Please ask for the correct username.\",\n",
    "            agent=authentication_agent,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# AGENTS\n",
    "\n",
    "# Human customer\n",
    "# Represents us as human-in-the-loop\n",
    "user = UserProxyAgent(\n",
    "    name=\"customer\",\n",
    "    code_execution_config=False,\n",
    ")\n",
    "\n",
    "# Triage Agent\n",
    "\n",
    "# Triage Agent prompt includes placeholders for the customer name, logged in status, and order ID\n",
    "order_triage_prompt = \"\"\"You are an order triage agent, working with a customer and a group of agents to provide support for your e-commerce platform.\n",
    "\n",
    "An agent needs to be logged in to be able to access their order. The authentication_agent will work with the customer to verify their identity, transfer to them to start with.\n",
    "The order_mgmt_agent will manage all order related tasks, such as tracking orders, managing orders, etc. Be sure to check the order as one step. Then if it's valid you can record it in the context.\n",
    "\n",
    "Use the check_order_id tool before the record_order_id tool, never together.\n",
    "\n",
    "Ask the customer for further information when necessary.\n",
    "\n",
    "The current status of this workflow is:\n",
    "Customer name: {customer_name}\n",
    "Logged in: {logged_in}\n",
    "Enquiring for Order ID: {order_id}\n",
    "\"\"\"\n",
    "\n",
    "order_triage_agent = ConversableAgent(\n",
    "    name=\"order_triage_agent\",\n",
    "    update_agent_state_before_reply=[\n",
    "        UpdateSystemMessage(order_triage_prompt),  # Inject the context variables into the system message\n",
    "    ],\n",
    "    functions=[check_order_id, record_order_id],  # Functions available to the agent\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "# Authentication Agent\n",
    "\n",
    "authentication_prompt = \"You are an authentication agent that verifies the identity of the customer.\"\n",
    "\n",
    "authentication_agent = ConversableAgent(\n",
    "    name=\"authentication_agent\",\n",
    "    system_message=authentication_prompt,\n",
    "    functions=[login_customer_by_username],  # Functions available to the agent\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "# Order Management Agent\n",
    "\n",
    "# Order Management Agent prompt includes placeholders for the customer name, logged in status, and order ID\n",
    "order_management_prompt = \"\"\"You are an order management agent that manages inquiries related to e-commerce orders.\n",
    "\n",
    "The order must be logged in to access their order.\n",
    "\n",
    "Use your available tools to get the status of the details from the customer. Ask the customer questions as needed.\n",
    "\n",
    "Use the check_order_id tool before the record_order_id tool, never together.\n",
    "\n",
    "The current status of this workflow is:\n",
    "Customer name: {customer_name}\n",
    "Logged in: {logged_in}\n",
    "Enquiring for Order ID: {order_id}\n",
    "\"\"\"\n",
    "\n",
    "order_mgmt_agent = ConversableAgent(\n",
    "    name=\"order_mgmt_agent\",\n",
    "    update_agent_state_before_reply=[\n",
    "        UpdateSystemMessage(order_management_prompt),  # Inject the context variables into the system message\n",
    "    ],\n",
    "    functions=[check_order_id, record_order_id],  # Functions available to the agent\n",
    "    llm_config=llm_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Nested Chats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NESTED CHAT - Delivery Status using a two-step chat queue:\n",
    "# The first step retrieves the order details with the order_retrieval_agent\n",
    "# The second step summarises the order details with the order_summariser_agent\n",
    "\n",
    "order_retrieval_agent = ConversableAgent(\n",
    "    name=\"order_retrieval_agent\",\n",
    "    system_message=\"You are an order retrieval agent that gets details about an order.\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "order_summariser_agent = ConversableAgent(\n",
    "    name=\"order_summariser_agent\",\n",
    "    system_message=\"You are an order summariser agent that provides a summary of the order details.\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "\n",
    "# Function to generate the message for the order_retireval_agent to use to output the order details\n",
    "def extract_order_summary(recipient: ConversableAgent, messages, sender: ConversableAgent, config):\n",
    "    \"\"\"Extracts the order summary based on the OrderID in the context variables\"\"\"\n",
    "    order_id = sender.get_context(\"order_id\")\n",
    "    if order_id in ORDER_DATABASE:\n",
    "        order = ORDER_DATABASE[order_id]\n",
    "        return f\"Order:\\n{json.dumps(order, indent=4)}\"\n",
    "    else:\n",
    "        return f\"Order {order_id} not found.\"\n",
    "\n",
    "\n",
    "nested_chat_one = {\n",
    "    \"carryover_config\": {\"summary_method\": \"last_msg\"},  # Bring the last message into the chat\n",
    "    \"recipient\": order_retrieval_agent,\n",
    "    \"message\": extract_order_summary,  # Retrieve the status details of the order using the order id\n",
    "    \"max_turns\": 1,  # Only one turn is necessary\n",
    "}\n",
    "\n",
    "nested_chat_two = {\n",
    "    \"recipient\": order_summariser_agent,\n",
    "    \"message\": \"Summarise the order details provided in a Markdown table format with headings: Detail, Information\",\n",
    "    \"max_turns\": 1,\n",
    "    \"summary_method\": \"last_msg\",  # Return their message back as the nested chat output\n",
    "}\n",
    "\n",
    "# Construct the chat queue\n",
    "chat_queue = [nested_chat_one, nested_chat_two]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Handoffs (OnCondition, OnContextCondition, and AfterWork)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HANDOFFS\n",
    "\n",
    "# Triage Agent handoffs\n",
    "register_hand_off(\n",
    "    agent=order_triage_agent,\n",
    "    hand_to=[\n",
    "        # If the customer is not logged in, transfer to the authentication agent\n",
    "        OnCondition(\n",
    "            target=authentication_agent,\n",
    "            condition=\"The customer is not logged in, authenticate the customer.\",\n",
    "            available=ContextExpression(\"!(${logged_in})\"),\n",
    "        ),\n",
    "        # If the customer is logged in, transfer to the order management agent\n",
    "        OnCondition(\n",
    "            target=order_mgmt_agent,\n",
    "            condition=\"The customer is logged in, continue with the order triage.\",\n",
    "            available=\"logged_in\",\n",
    "        ),\n",
    "        # If the order enquiry has been fulfilled, transfer back to the user\n",
    "        AfterWork(AfterWorkOption.REVERT_TO_USER),\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Authentication Agent handoffs\n",
    "register_hand_off(\n",
    "    agent=authentication_agent,\n",
    "    hand_to=[\n",
    "        # When logged in, transfer to the order triage agent (no LLM required)\n",
    "        # OnContextCondition checks the context variables and handsoff immediately\n",
    "        OnContextCondition(\n",
    "            target=order_triage_agent,\n",
    "            condition=\"logged_in\",\n",
    "        ),\n",
    "        # If the customer cannot be authenticated, tell the user so it can try again\n",
    "        AfterWork(AfterWorkOption.REVERT_TO_USER),\n",
    "    ],\n",
    ")\n",
    "\n",
    "\n",
    "# Function to check if we have an Order ID, used in the Order Management Agent handoff\n",
    "def has_order_in_context(agent: ConversableAgent, messages: list[dict[str, Any]]) -> bool:\n",
    "    return agent.get_context(\"has_order_id\")\n",
    "\n",
    "\n",
    "# Order Management Agent handoffs\n",
    "register_hand_off(\n",
    "    agent=order_mgmt_agent,\n",
    "    hand_to=[\n",
    "        # If we have an Order ID, transfer to the nested chat to get the order summary\n",
    "        OnCondition(\n",
    "            target={\n",
    "                \"chat_queue\": chat_queue,\n",
    "            },\n",
    "            condition=\"Retrieve the status of the order\",\n",
    "            available=has_order_in_context,\n",
    "        ),\n",
    "        # If the user is not logged in, transfer to the authentication agent\n",
    "        OnCondition(\n",
    "            target=authentication_agent,\n",
    "            condition=\"The customer is not logged in, authenticate the customer.\",\n",
    "            available=ContextExpression(\"!(${logged_in})\"),\n",
    "        ),\n",
    "        # If the customer has no more enquiries about this order, transfer back to the\n",
    "        # Triage agent to handle it enquiry\n",
    "        OnCondition(target=order_triage_agent, condition=\"The customer has no more enquiries about this order.\"),\n",
    "        # As a fallback, transfer back to the user for more information\n",
    "        AfterWork(AfterWorkOption.REVERT_TO_USER),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's go!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the chat with the Triage Agent\n",
    "chat_history = initiate_swarm_chat(\n",
    "    initial_agent=order_triage_agent,\n",
    "    agents=[order_triage_agent, authentication_agent, order_mgmt_agent],\n",
    "    context_variables=workflow_context,  # Our shared context\n",
    "    messages=\"Can you help me with my order.\",  # Starting message\n",
    "    user_agent=user,  # Human-in-the-loop\n",
    "    max_rounds=40,  # Maximum number of turns\n",
    "    after_work=AfterWorkOption.TERMINATE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Controlling flow\n",
    "\n",
    "### When not logged in\n",
    "![Swarm Enhanced Demonstration](https://media.githubusercontent.com/media/ag2ai/ag2/refs/heads/main/notebook/swarm_enhanced_02.png)\n",
    "\n",
    "### When logged in but no order id\n",
    "![Swarm Enhanced Demonstration](https://media.githubusercontent.com/media/ag2ai/ag2/refs/heads/main/notebook/swarm_enhanced_03.png)\n",
    "\n",
    "### When logged in with order id\n",
    "![Swarm Enhanced Demonstration](https://media.githubusercontent.com/media/ag2ai/ag2/refs/heads/main/notebook/swarm_enhanced_04.png)\n",
    "\n",
    "# Agent state\n",
    "\n",
    "### Agent System Messages with context\n",
    "![Swarm Enhanced Demonstration](https://media.githubusercontent.com/media/ag2ai/ag2/refs/heads/main/notebook/swarm_enhanced_05.png)"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Swarm Ochestration",
   "tags": [
    "orchestration",
    "group chat",
    "swarm"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
